package org.misty.redis.locker;

import lombok.Getter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.misty.redis.client.RedisClient;
import org.misty.redis.util.Util;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.function.Function;

import static org.misty.redis.util.Util.checkArgument;

/**
 * @author Misty on 2020-05-31
 */
@Slf4j
@ToString(onlyExplicitlyIncluded = true)
public class RedisLocker {
    private final static ThreadLocal<Map<Object, Context>> ctxMapLocal = ThreadLocal.withInitial(HashMap::new);
    private final static String ACQUIRE_LUA;
    private final static String RELEASE_LUA;

    public final static int ACQUIRE_NEVER_TIMEOUT = 0;

    public final static int DEFAULT_VALUE_EXPIRE = 10000;
    public final static int DEFAULT_VALUE_ACQUIRE = ACQUIRE_NEVER_TIMEOUT;

    static {
        ACQUIRE_LUA = Util.copyToString(RedisLocker.class.getClassLoader().getResourceAsStream("lua/locker/acquire.lua"));
        RELEASE_LUA = Util.copyToString(RedisLocker.class.getClassLoader().getResourceAsStream("lua/locker/release.lua"));
    }

    @Getter
    private final RedisClient client;

    @ToString.Include
    private final String hash = Integer.toHexString(System.identityHashCode(this));
    @ToString.Include
    private final String host;
    @ToString.Include
    private final int port;
    @Getter
    @ToString.Include
    private final int expireTimeMills;
    @Getter
    @ToString.Include
    private final int acquireTimeoutMills;

    public RedisLocker(RedisClient redisClient, ILockerConfig config) {
        Objects.requireNonNull(redisClient);
        Objects.requireNonNull(config);

        this.expireTimeMills = Optional.ofNullable(config.getExpireTime()).orElse(DEFAULT_VALUE_EXPIRE);
        checkArgument(expireTimeMills > 0, "Expire time in milliseconds expireTime(%s) must be > 0", expireTimeMills);

        this.acquireTimeoutMills = Optional.ofNullable(config.getAcquireTimeout()).orElse(DEFAULT_VALUE_ACQUIRE);
        checkArgument(acquireTimeoutMills >= 0, "Acquire timeout in milliseconds acquireTimeout(%s) must be >= 0", acquireTimeoutMills);

        this.client = redisClient;
        this.host = redisClient.getHost();
        this.port = redisClient.getPort();

        log.info("Initialized with {expireTime={}, acquireTimeout={}}", expireTimeMills, acquireTimeoutMills);
    }

    /**
     * 获得锁
     * 使用默认的过期时间和请求超时时间
     */
    public String acquire(String key, String token) {
        Objects.requireNonNull(key);
        Objects.requireNonNull(token);
        return acquireInternal(key, token, expireTimeMills, acquireTimeoutMills, false);
    }

    /**
     * 获得锁
     * 使用指定的过期时间，默认的请求超时时间
     */
    public String acquire(String key, String token, int expireTime) {
        return acquireInternal(key, token, expireTime, ACQUIRE_NEVER_TIMEOUT, true);
    }

    /**
     * 获得锁
     * 使用指定的过期时间和请求超时时间
     */
    public String acquire(String key, String token, int expireTime, int acquireTimeout) {
        return acquireInternal(key, token, expireTime, acquireTimeout, true);
    }

    private String acquireInternal(String key, String token, int expireTime, int acquireTimeout, boolean checkArgs) {
        if (checkArgs) {
            Objects.requireNonNull(key);
            Objects.requireNonNull(token);
            checkArgument(expireTime > 0, "Expire time in milliseconds expireTime(%s) must be > 0", expireTime);
            checkArgument(acquireTimeout >= 0, "Acquire timeout in milliseconds acquireTimeout(%s) must be >= 0", acquireTimeout);
        }

        Thread t = Thread.currentThread();
        long end = System.currentTimeMillis() + acquireTimeout;

        List<String> keys = Collections.singletonList(key);
        String[] values = {token, String.valueOf(expireTime)};

        try {
            while ((acquireTimeout <= ACQUIRE_NEVER_TIMEOUT || System.currentTimeMillis() < end) && !t.isInterrupted()) {
                Long result = client.eval(ACQUIRE_LUA, Long.class, keys, values);
                if (1L == result) {
                    log.debug("acquire lock success, key is {}, token is {}", key, token);
                    return token;
                }
                try {
                    TimeUnit.MILLISECONDS.sleep(50);
                } catch (InterruptedException e) {
                    t.interrupt();
                }
            }
        } catch (Exception e) {
            log.error("acquire lock due to error", e);
        }
        return null;
    }

    /**
     * 尝试获得锁
     * 使用默认的过期时间
     */
    public String tryAcquire(String key, String token) {
        Objects.requireNonNull(key);
        Objects.requireNonNull(token);
        return tryAcquireInternal(key, token, expireTimeMills, false);
    }

    /**
     * 尝试获得锁
     * 使用指定的过期时间
     */
    public String tryAcquire(String key, String token, int expireTime) {
        return tryAcquireInternal(key, token, expireTime, true);
    }

    private String tryAcquireInternal(String key, String token, int expireTime, boolean checkArgs) {
        if (checkArgs) {
            Objects.requireNonNull(key);
            Objects.requireNonNull(token);
            checkArgument(expireTime > 0, "Expire time in milliseconds expireTime(%s) must be > 0", expireTime);
        }
        List<String> keys = Collections.singletonList(key);
        String[] values = {token, String.valueOf(expireTime)};
        try {
            Long result = client.eval(ACQUIRE_LUA, Long.class, keys, values);
            if (1L == result) {
                log.debug("acquire lock success, key is {}, token is {}", key, token);
                return token;
            }
        } catch (Exception e) {
            log.error("acquire lock due to error", e);
        }
        return null;
    }

    public boolean release(String key, String token) {
        Objects.requireNonNull(token);

        List<String> keys = Collections.singletonList(key);
        String[] values = {token};
        try {
            Long result = client.eval(RELEASE_LUA, Long.class, keys, values);
            if (1L == result) {
                log.debug("release lock success, token is {}", token);
                return true;
            }
        } catch (Exception e) {
            log.error("release lock due to error", e);
        }
        return false;
    }

    public void lock(final String key, final Consumer<Context> func) {
        this.lock(key, ctx -> {
            func.accept(ctx);
            return null;
        });
    }

    public <R> R lock(final String key, final Function<Context, R> func) {
        Objects.requireNonNull(key);
        Objects.requireNonNull(func);

        String token;
        Map<Object, Context> ctxMap = ctxMapLocal.get();
        Context ctx = ctxMap.get(key);
        if (ctx == null) {
            token = UUID.randomUUID().toString();
            ctx = new Context(key, token);
            ctxMap.put(key, ctx);
        } else {
            token = ctx.getToken();
        }
        ctx.numLock++;

        try {
            String acquireToken = acquire(key, token);
            if (token.equals(acquireToken)) {
                try {
                    return func.apply(ctx);
                } finally {
                    release(key, token);
                }
            }
        } finally {
            ctx.numLock--;
            if (ctx.numLock == 0) {
                ctxMap.remove(key);
            }
            if (ctxMap.isEmpty()) {
                ctxMapLocal.remove();
            }
        }
        return null;
    }
}
